package com.zlz.m3u8download.service.impl;

import com.baomidou.mybatisplus.core.toolkit.AES;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.zlz.m3u8download.bean.RunTimeTask;
import com.zlz.m3u8download.bean.TaskCacheResult;
import com.zlz.m3u8download.bean.entity.CacheM3u8Entity;
import com.zlz.m3u8download.bean.entity.M3u8HLCFileEntity;
import com.zlz.m3u8download.bean.enums.CacheM3u8State;
import com.zlz.m3u8download.bean.enums.DownloadStatus;
import com.zlz.m3u8download.bean.enums.M3u8Type;
import com.zlz.m3u8download.config.M3u8Config;
import com.zlz.m3u8download.mapper.CacheM3u8Mapper;
import com.zlz.m3u8download.service.TaskService;
import com.zlz.m3u8download.task.CommonCache;
import com.zlz.m3u8download.util.FileUtil;
import com.zlz.m3u8download.util.M3u8FileUtil;
import com.zlz.m3u8download.util.OkHttpUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @Author: zhulinzhong
 * @Date: 2024/11/29
 */
@Service
@Slf4j
public class TaskServiceImpl implements TaskService {

    private final CacheM3u8Mapper cacheM3u8Mapper;

    private final M3u8Config config;

    public TaskServiceImpl(CacheM3u8Mapper cacheM3u8Mapper,
                           M3u8Config config) {
        this.cacheM3u8Mapper = cacheM3u8Mapper;
        this.config = config;
    }

    private List<M3u8HLCFileEntity> createFragmentsTask(CacheM3u8Entity list, CacheM3u8Entity video) throws IOException {

        // 生成video列表
        List<String> rows = M3u8FileUtil.split(video.getContent());
        List<String> sourceRows = M3u8FileUtil.split(video.getSourceContent());
        List<M3u8HLCFileEntity> fragments = new ArrayList<>();
        int start = 10001;
        String method = null;
        String key = null;
        String iv = null;
        for (int i = 0; i < sourceRows.size(); i++) {
            String row = sourceRows.get(i);
            if (row.startsWith("#")) {
                // 检查是否加密
                if(row.startsWith("#EXT-X-KEY")){
                    String[] split = row.split(":");
                    String[] kvs = split[1].split(",");
                    Map<String, String> kvMap = new HashMap<>();
                    for (String kvStr : kvs) {
                        String[] kv = kvStr.split("=");
                        kvMap.put(kv[0], kv[1]);
                    }
                    String url = FileUtil.buildUrl(list.getSourceUrl(), kvMap.get("URI").substring(1, kvMap.get("URI").length() - 1));
                    method = kvMap.get("METHOD");
                    key = OkHttpUtil.get(url);
                    iv = kvMap.get("IV");
                }
                continue;
            }

            M3u8HLCFileEntity fragment = M3u8HLCFileEntity.builder()
                    .code(start++)
                    .parentId(video.getId())
                    .downloadUrl(FileUtil.buildUrl(list.getSourceUrl(), row))
                    .filename(rows.get(i).replace(config.domainUrl, ""))
                    .fileType(config.suffix)
                    .state(DownloadStatus.PENDING_DOWNLOAD.getValue())
                    .tryTimes(0)
                    .method(method)
                    .key(key)
                    .iv(iv)
                    .build();
            fragments.add(fragment);
        }
        return fragments;
    }


    @Override
    public TaskCacheResult startTaskForPlay(CacheM3u8Entity cacheM3u8, Long userId) {
        TaskCacheResult result = TaskCacheResult.builder().build();
        List<CacheM3u8Entity> videoTask = new ArrayList<>();
        if (cacheM3u8.getType().equals(M3u8Type.VIDEO.getValue())) {
            videoTask.add(cacheM3u8);
        } else {
            List<CacheM3u8Entity> cacheM3u8Entities = cacheM3u8Mapper.selectList(
                    Wrappers.<CacheM3u8Entity>lambdaQuery().eq(CacheM3u8Entity::getParentId, cacheM3u8.getId())
            );
            videoTask.addAll(cacheM3u8Entities);
        }
        CommonCache.fastStart(videoTask.get(0).getId(), userId);
        result.setTaskId(String.valueOf(videoTask.get(0).getId()));
        result.setUrl(cacheM3u8.getUrl());
        return result;
    }

    @Override
    public String stopTask(Long taskId) {
        CommonCache.priorityTask.remove(taskId);
        CommonCache.waitingTasks.remove(taskId);
        RunTimeTask runTimeTask = CommonCache.runningTasks.get(taskId);
        if (runTimeTask != null) {
            runTimeTask.setRunning(false);
        }
        return "任务已取消";
    }

    @Override
    public boolean execute(RunTimeTask runTimeTask) {
        // 1、获取视频信息
        CacheM3u8Entity task = cacheM3u8Mapper.selectOne(Wrappers.<CacheM3u8Entity>lambdaQuery()
                .eq(CacheM3u8Entity::getId, runTimeTask.getTaskId()));
        // 2、检查任务状态
        if (task == null) {
            return true;
        }
        if (CacheM3u8State.CACHED_COMPLETED.getValue().equals(task.getState())) {
            return true;
        }

        log.debug("正在执行任务,TaskId:{}", runTimeTask.getTaskId());
        try {
            // list索引落盘
            CacheM3u8Entity list = cacheM3u8Mapper.selectOne(Wrappers.<CacheM3u8Entity>lambdaQuery()
                    .eq(CacheM3u8Entity::getId, task.getParentId()));
            if (list != null) {
                FileUtil.save(config.localPath, list.getPath() + config.m3u8Filename, list.getContent().getBytes(StandardCharsets.UTF_8));
            }

            // video索引落盘
            FileUtil.save(config.localPath, task.getPath() + config.m3u8Filename, task.getContent().getBytes(StandardCharsets.UTF_8));
            // 任务切片
            List<M3u8HLCFileEntity> fragments = createFragmentsTask(list, task);
            log.debug("任务切片完成,任务Id:{},总数量:{}", runTimeTask.getTaskId(), fragments.size());

            // 创建多线程下载
            ThreadPoolExecutor threadPoolExecutor = new ThreadPoolExecutor(
                    10,
                    10,
                    60,
                    TimeUnit.SECONDS,
                    new LinkedBlockingQueue<>());
            runTimeTask.setFragments(fragments.stream().collect(Collectors.toMap(M3u8HLCFileEntity::getCode, Function.identity())));
            runTimeTask.setMaxCode(fragments.get(fragments.size() - 1).getCode());
            while (runTimeTask.isRunning()) {
                if (runTimeTask.getRunningTaskCode().size() > 10) {
                    Thread.sleep(20);
                    continue;
                }
                if (runTimeTask.getWaitingTaskCode().isEmpty()) {
                    break;
                }
                Set<Integer> waitingTaskCode = runTimeTask.getWaitingTaskCode();
                Integer code = waitingTaskCode.iterator().next();
                runTimeTask.updateFailedTask(code, false);
                runTimeTask.updateRunningTask(code, true);
                threadPoolExecutor.execute(() -> saveFile(runTimeTask.getFragments().get(code), runTimeTask));
            }

            // 主动停止
            if (!runTimeTask.isRunning()) {
                log.debug("任务主动终止:任务:{}", task.getId());
                threadPoolExecutor.shutdown();
                log.debug("任务终止成功:任务:{}", task.getId());
                return true;
            }

            // 运行完成
            long count = fragments.stream().filter((fragment) -> !DownloadStatus.DOWNLOAD_SUCCESS.getValue().equals(fragment.getCode())).count();
            if (count != fragments.size()) {
                log.info("任务执行部分成功:任务:{}", task.getId());
                return false;
            }

            if (list != null) {
                list.setState(CacheM3u8State.CACHED_COMPLETED.getValue());
                cacheM3u8Mapper.updateById(list);
            }
            task.setState(CacheM3u8State.CACHED_COMPLETED.getValue());
            cacheM3u8Mapper.updateById(task);
            threadPoolExecutor.shutdown();
            return true;
        } catch (Exception e) {
            log.error("执行任务失败", e);
            return false;
        }
    }

    @Override
    public void taskStatus(Long taskId) {

    }

    public void saveFile(M3u8HLCFileEntity fragment, RunTimeTask runTimeTask) {
        log.debug("切片任务执行中:TaskId:{}-{}", fragment.getParentId(), fragment.getCode());
        try {
            if (DownloadStatus.DOWNLOAD_SUCCESS.getValue().equals(fragment.getState())) {
                return;
            }
            if (FileUtil.contains(config.localPath, fragment.getFilename())) {
                fragment.setState(DownloadStatus.DOWNLOAD_SUCCESS.getValue());
                runTimeTask.updateRunningTask(fragment.getCode(), false);
                log.debug("切片已存在，直接返回成功:TaskId:{}-{}", fragment.getParentId(), fragment.getCode());
                return;
            }

            // 执行下载任务
            byte[] data = OkHttpUtil.getBytes(fragment.getDownloadUrl());

            // 解密
            if(fragment.getMethod() != null){
                data = AES.decrypt(data, fragment.getKey().getBytes());
            }

            // 保存到本地
            if (!FileUtil.save(config.localPath, fragment.getFilename(), data)) {
                runTimeTask.updateRunningTask(fragment.getCode(), false);
                runTimeTask.updateFailedTask(fragment.getCode(), true);
                log.error("切片任务执行失败:TaskId:{}-{}", fragment.getParentId(), fragment.getCode());
                return;
            }
            // 更新任务状态
            fragment.setState(DownloadStatus.DOWNLOAD_SUCCESS.getValue());
            runTimeTask.updateRunningTask(fragment.getCode(), false);
            log.debug("切片任务执行成功:TaskId:{}-{}", fragment.getParentId(), fragment.getCode());
        } catch (Exception e) {
            runTimeTask.updateRunningTask(fragment.getCode(), false);
            runTimeTask.updateFailedTask(fragment.getCode(), true);
            log.error("切片任务执行失败:TaskId:{}-{}", fragment.getParentId(), fragment.getCode());
            log.error("文件保存异常", e);
        }
    }

}
